How to decorate a function that is used for a curve_fit